import numpy as np
from scipy.ndimage import affine_transform
from components.VolumeInfo import VolumeInfo


def _fast_apply_window(hu_array, window_center, window_width):
    min_value = window_center - window_width // 2
    max_value = window_center + window_width // 2
    hu_array = np.clip(hu_array, min_value, max_value)
    normalized = ((hu_array - min_value) / window_width * 255.0)
    return normalized


def _fast_transfer_to_hu(pixel_array, rescale_slope, rescale_intercept):
    return pixel_array * rescale_slope + rescale_intercept


class BaseImgLayer:
    def __init__(self, view_type, dicom_tag_info, volume_info):
        self.dicom_tag_info = dicom_tag_info
        self.volume_info: VolumeInfo = volume_info
        self.hu_array = None
        self.view_type = view_type

    def get_base_img_arr(self, hu_array):
        hu_array = np.asarray(hu_array, dtype=np.float32)
        self.hu_array = hu_array
        window_center, window_width = self.volume_info.view_info[self.view_type].window

        pixel_array = _fast_apply_window(hu_array, window_center, window_width)
        transform_pixel_array = self.do_transfer(pixel_array)
        pixel_array_8bit = transform_pixel_array.astype(np.uint8)
        return pixel_array_8bit

    # stack视图大约需要0.01秒。
    def do_transfer(self, pixel_array):
        matrix = self.volume_info.view_info[self.view_type].transform_matrix
        inverse_matrix = np.linalg.inv(matrix)
        # 只需要3x3矩阵的前两行
        affine_matrix = inverse_matrix[:2, :2]
        offset = inverse_matrix[:2, 2]
        output_shape = self.volume_info.view_info[self.view_type].png_size
        # 执行变换
        transformed = affine_transform(
            pixel_array,
            affine_matrix,
            offset=offset,
            output_shape=output_shape,
            order=1
        )
        return transformed

    def get_hu_value_by_position(self, dicom_position):
        return self.hu_array[dicom_position[1]][dicom_position[0]]

